import mediapipe as mp
from cv2 import cv2

mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose

path = ".\\pose.mp4"

from tornado import web, httpserver, ioloop, websocket
import os
import cv2
import math
import numpy as np
from PIL import Image
import base64

port = 10101
address = "localhost"
JPEG_HEADER = "data:image/jpeg;base64,"

# TODO: 摄像头就把path换成0
cap = cv2.VideoCapture(0)
if not cap.isOpened():
    print('quit')
    quit()
ret, frame = cap.read()

pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5)


def calculate_angle(point_1, point_2, point_3):
    point_1 = np.array(point_1)
    point_2 = np.array(point_2)
    point_3 = np.array(point_3)

    a = math.sqrt(
        (point_2[0] - point_3[0]) * (point_2[0] - point_3[0])
        + (point_2[1] - point_3[1]) * (point_2[1] - point_3[1])
        + (point_2[2] - point_3[2]) * (point_2[2] - point_3[2]))
    b = math.sqrt(
        (point_1[0] - point_3[0]) * (point_1[0] - point_3[0])
        + (point_1[1] - point_3[1]) * (point_1[1] - point_3[1])
        + (point_1[2] - point_3[2]) * (point_1[2] - point_3[2]))
    c = math.sqrt(
        (point_1[0] - point_2[0]) * (point_1[0] - point_2[0])
        + (point_1[1] - point_2[1]) * (point_1[1] - point_2[1])
        + (point_1[2] - point_2[2]) * (point_1[2] - point_2[2]))
    angle = math.degrees(math.acos((b * b - a * a - c * c) / (-2 * a * c)))
    if angle > 180:
        angle = 360 - angle
    return angle


def calculate_horizon_angle(point_1, point_2, point_3):
    point_1 = np.array(point_1)
    point_2 = np.array(point_2)
    point_3 = np.array(point_3)

    a = math.sqrt(
        (point_2[0] - point_3[0]) * (point_2[0] - point_3[0]) + (point_2[2] - point_3[2]) * (point_2[2] - point_3[2]))
    b = math.sqrt(
        (point_1[0] - point_3[0]) * (point_1[0] - point_3[0]) + (point_1[2] - point_3[2]) * (point_1[2] - point_3[2]))
    c = math.sqrt(
        (point_1[0] - point_2[0]) * (point_1[0] - point_2[0]) + (point_1[2] - point_2[2]) * (point_1[2] - point_2[2]))
    angle = math.degrees(math.acos((b * b - a * a - c * c) / (-2 * a * c)))
    if angle > 180:
        angle = 360 - angle
    return angle

    # radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
    # angle = np.abs(radians * 180.0 / np.pi)
    #
    # if angle > 180.0:
    #     angle = 360 - angle
    # a = np.array(a)
    # b = np.array(b)
    # c = np.array(c)
    #
    # mid_a = a - b
    # mid_c = c - b
    #
    # dot = mid_a.dot(mid_c)
    #
    # l_a = np.sqrt(mid_a.dot(mid_a))
    # l_c = np.sqrt(mid_c.dot(mid_c))
    #
    # cos_a_c = dot/(l_a*l_c)
    #
    # angle = np.arccos(cos_a_c)*180/np.pi
    #
    # radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
    # angle = np.abs(radians * 180.0 / np.pi)
    #
    # if angle > 180.0:
    #     angle = 360 - angle
    #
    # return angle


def getAngle(results, F, M, E):
    landmarks = results.pose_world_landmarks.landmark

    point_1 = [landmarks[F.value].x, landmarks[F.value].y, landmarks[F.value].z]
    point_2 = [landmarks[M.value].x, landmarks[M.value].y, landmarks[M.value].z]
    point_3 = [landmarks[E.value].x, landmarks[E.value].y, landmarks[E.value].z]  # mp_pose.PoseLandmark.LEFT_WRIST

    angle = calculate_angle(point_1, point_2, point_3)
    return angle


def getHorizonAngle(results, F, M, E):
    landmarks = results.pose_world_landmarks.landmark

    point_1 = [landmarks[F.value].x, landmarks[F.value].y, landmarks[F.value].z]
    point_2 = [landmarks[M.value].x, landmarks[M.value].y, landmarks[M.value].z]
    point_3 = [landmarks[E.value].x, landmarks[E.value].y, landmarks[E.value].z]  # mp_pose.PoseLandmark.LEFT_WRIST

    angle = calculate_horizon_angle(point_1, point_2, point_3)
    return angle


def getData(results):
    left_armpit = getAngle(results, mp_pose.PoseLandmark.LEFT_ELBOW,
                           mp_pose.PoseLandmark.LEFT_SHOULDER,
                           mp_pose.PoseLandmark.LEFT_HIP)
    left_arm_shoulder_horizon = getHorizonAngle(results, mp_pose.PoseLandmark.LEFT_ELBOW,
                                                mp_pose.PoseLandmark.LEFT_SHOULDER,
                                                mp_pose.PoseLandmark.RIGHT_SHOULDER)
    right_armpit = getAngle(results, mp_pose.PoseLandmark.RIGHT_ELBOW,
                            mp_pose.PoseLandmark.RIGHT_SHOULDER,
                            mp_pose.PoseLandmark.RIGHT_HIP)
    right_arm_shoulder_horizon = getHorizonAngle(results, mp_pose.PoseLandmark.RIGHT_ELBOW,
                                                 mp_pose.PoseLandmark.RIGHT_SHOULDER,
                                                 mp_pose.PoseLandmark.LEFT_SHOULDER)
    pose1 = {'right_armpit': right_armpit, 'right_arm_shoulder_horizon': right_arm_shoulder_horizon,
             'left_armpit': left_armpit, 'left_arm_shoulder_horizon': left_arm_shoulder_horizon}
    return pose1


def get_image_dataurl():
    ret, frame = cap.read()
    if ret == False or frame is None:
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)

    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    image.flags.writeable = False

    results = pose.process(image)

    image.flags.writeable = True
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
                              mp_drawing.DrawingSpec(color=(245, 117, 66), thickness=2, circle_radius=2),
                              mp_drawing.DrawingSpec(color=(245, 66, 230), thickness=2, circle_radius=2)
                              )

    r, buf = cv2.imencode(".jpeg", image)
    dat = Image.fromarray(np.uint8(buf)).tobytes()
    img_date_url = JPEG_HEADER + str(base64.b64encode(dat))[2:-1]
    data = getData(results)
    return img_date_url, data


class VideoHandler(websocket.WebSocketHandler):
    def check_origin(self, origin):
        return True

    def on_message(self, message):
        img, data = get_image_dataurl()
        self.write_message({"img": img, "data": data})


app = web.Application(handlers=[(r'/video', VideoHandler)],
                      template_path=os.path.join(os.path.dirname(__file__), "ui"))
http_server = httpserver.HTTPServer(app)
http_server.listen(port=port, address=address)
print("URL: http://{}:{}/index".format(address, port))
ioloop.IOLoop.instance().start()
